bl_info = {
    "name": "LegoColorChange",
    "description": "dae files from Stuido2.0(bricklink)lost color information, this add-on is for solving this problem.",
    "author": "zju06",
    "version": (0, 0, 10),
    "blender": (2, 91, 2),
    "location": "3D View > Tools",
    "warning": "", # used for warning icon and text in addons panel
    "wiki_url": "",
    "tracker_url": "",
    "category": "Development"
}


import bpy
import time

from bpy.props import (StringProperty,
                       BoolProperty,
                       IntProperty,
                       FloatProperty,
                       FloatVectorProperty,
                       EnumProperty,
                       PointerProperty,
                       )
from bpy.types import (Panel,
                       Menu,
                       Operator,
                       PropertyGroup,
                       )
model_loc = [0.0,0.0,0.0]
# ------------------------------------------------------------------------
#    lego color dict
# ------------------------------------------------------------------------ 
lego_color_dict = {
'SOLID-WHITE':[1,1,1,1],
'SOLID-LIGHT_GRAY':[0.565,0.578,0.546,1],
'SOLID-DARK_GRAY':[0.25,0.242,0.22,1],
'SOLID-LIGHT_BLUISH_GRAY':[0.497,0.521,0.527,1],
'SOLID-DARK_BLUISH_GRAY':[0.15,0.175,0.175,1],
'SOLID-BLACK':[0,0,0,1],
'SOLID-BLUE':[0.01,0.01,1,1],
'SOLID-LIGHT_BLUE':[0.45,0.65,0.9,1],
'SOLID-BRIGHT_LIGHT_BLUE':[0.3,0.9,1,1],
'SOLID-SKY_BLUE':[0.05,0.75,0.9,1],
'SOLID-MAERSK_BLUE':[0.242,0.57,0.7,1],
'SOLID-SAND_BLUE':[0.15,0.205,0.271,1],
'SOLID-DARK_BLUE':[0.035,0.051,0.105,1],
'SOLID-VIOLET':[0.065,0.065,0.356,1],
'SOLID-LIGHT_VIOLET':[0.8,0.95,1,1],
'SOLID-MEDIUM_VIOLET':[0.287,0.366,0.95,1],
'SOLID-BLUE_VIOLET':[0.1,0.175,0.5,1],
'SOLID-DARK_BLUE_VIOLET':[0.05,0.05,1,1],
'SOLID-LIGHT_TURQUOISE':[0.08,0.58,0.48,1],
'SOLID-DARK_TURQUOISE':[0.02,0.37,0.32,1],
'SOLID-MEDIUM_AZURE':[0.1,0.42,0.82,1],
'SOLID-DARK_AZURE':[0.02,0.25,0.6,1],
'SOLID-GREEN':[0.02,0.32,0.117,1],
'SOLID-YELLOWISH_GREEN':[1,1,0.214,1],
'SOLID-LIGHT_GREEN':[0.409,1,0.287,1],
'SOLID-MEDIUM_GREEN':[0.205,1,0.345,1],
'SOLID-NEON_GREEN':[0.884,1,0.116,1],
'SOLID-LIME':[0.365,1,0.037,1],
'SOLID-LIGHT_LIME':[0.943,1,0,1],
'SOLID-MEDIUM_LIME':[0.471,1,0,1],
'SOLID-SAND_GREEN':[0.231,0.361,0.246,1],
'SOLID-BRIGHT_GREEN':[0.054,0.418,0.102,1],
'SOLID-OLIVE_GREEN':[0.275,0.275,0.095,1],
'SOLID-DARK_GREEN':[0,0.2,0.2,1],
'SOLID_AQUA':[0.8,0.8,0.8,1],
'SOLID_LIGHT_AQUA':[0.8,0.8,0.8,1],
'SOLID-YELLOW':[1,1,0,1],
'SOLID-BRIGHT_LIGHT_YELLOW':[1,1,0.143,1],
'SOLID-LIGHT_YELLOW':[1,1,0.252,1],
'SOLID-DARK_YELLOW':[1,0.429,0.045,1],
'SOLID-ORANGE':[1,0.15,0,1],
'SOLID-VERY_LIGHT_ORANGE':[1,0.638,0.179,1],
'SOLID-BRIGHT_LIGHT_ORANGE':[1,0.49,0.035,1],
'SOLID-LIGHT_ORANGE':[1,0.313,0.038,1],
'SOLID-MEDIUM_ORANGE':[1,0.29,0.038,1],
'SOLID-EARTH_ORANGE':[1,0.213,0.034,1],
'SOLID-FABULAND_ORANGE':[0.9,0.2,0.025,1],
'SOLID-NEON_ORANGE':[1,0.212,0.126,1],
'SOLID-DARK_ORANGE':[1,0.061,0,1],
'SOLID-NOUGAT':[1,0.413,0.202,1],
'SOLID-LIGHT_NOUGAT':[1,0.784,0.339,1],
'SOLID-MEDIUM_NOUGAT':[1,0.272,0.114,1],
'SOLID-DARK_NOUGAT':[1,0.205,0.067,1],
'SOLID-TAN':[0.8,0.7,0.4,1],
'SOLID-DARK_TAN':[0.785,0.5,0.2,1],
'SOLID-RUST':[0.421,0,0,1],
'SOLID-RED':[1,0,0,1],
'SOLID-SAND_RED':[0.532,0.188,0.196,1],
'SOLID-DARK_RED':[0.4,0,0,1],
'SOLID-BROWN':[0.2,0.067,0.045,1],
'SOLID-FABULAND_BROWN':[0.61,0.216,0.109,1],
'SOLID-MEDIUM_BROWN':[0.337,0.133,0.051,1],
'SOLID-REDDISH_BROWN':[0.202,0.072,0.044,1],
'SOLID-DARK_BROWN':[0.114,0.08,0.067,1],
'SOLID-PINK':[1,0.5,0.5,1],
'SOLID-LIGHT_PINK':[1,0.629,0.721,1],
'SOLID-BRIGHT_PINK':[1,0.6,0.6,1],
'SOLID-MEDIUM_DARK_PINK':[0.788,0.24,0.24,1],
'SOLID-DARK_PINK':[0.736,0.148,0.204,1],
'SOLID-PURPLE':[0.227,0.013,0.332,1],
'SOLID-LIGHT_PURPLE':[0.397,0.074,0.446,1],
'SOLID-SAND_PURPLE':[0.38,0.239,0.466,1],
'SOLID-DARK_PURPLE':[0.122,0.076,0.305,1],
'SOLID-MAGENTA':[0.527,0.068,0.262,1],
'SOLID-LAVENDER':[0.773,0.489,1,1],
'SOLID-MEDIUM_LAVENDER':[0.489,0.167,0.708,1],
'SOLID-LIGHT_SALMON':[1,0.568,0.568,1],
'SOLID-SALMON':[1,0.164,0.164,1],
'SOLID-CORAL':[1,0.121,0.121,1],
}                      

# ------------------------------------------------------------------------
#    Scene Properties
# ------------------------------------------------------------------------

class MyProperties(PropertyGroup):

    my_bool: BoolProperty(
        name="Enable or Disable",
        description="A bool property",
        default = False
        )

    my_int: IntProperty(
        name = "Int Value",
        description="A integer property",
        default = 23,
        min = 10,
        max = 100
        )

    my_float: FloatProperty(
        name = "Float Value",
        description = "A float property",
        default = 23.7,
        min = 0.01,
        max = 30.0
        )

    my_float_vector: FloatVectorProperty(
        name = "Float Vector Value",
        description="Something",
        default=(0.0, 0.0, 0.0), 
        min= 0.0, # float
        max = 0.1
    ) 

    my_string: StringProperty(
        name="User Input",
        description=":",
        default="",
        maxlen=1024,
        )

    my_path: StringProperty(
        name = "Sheet",
        description="Choose the color_sheet:",
        default="",
        maxlen=1024,
        subtype='DIR_PATH'
        )

    my_enum: EnumProperty(
        name="Dropdown:",
        description="Apply Data to attribute.",
        items=[ ('OP1', "Option 1", ""),
                ('OP2', "Option 2", ""),
                ('OP3', "Option 3", ""),
               ]
        )

# ------------------------------------------------------------------------
#    Operators
# ------------------------------------------------------------------------

    
class WM_OT_LegoDaeCorrect(Operator):
    bl_label = "Quick Correct"
    bl_idname = "wm.lego_dae_correct"
    
    def execute(self, context):
        scene = context.scene
        mytool = scene.my_tool
        
        lego_waste_list = (
            "Light", "Light_Node_1","Light_Node_2",
            "Camera_Node","Camera","Cube")
        # delete cameras, lights by default and from Studio2.0
        for i in lego_waste_list:
            node = bpy.data.objects.get(i)
            if node != None:
                bpy.data.objects.remove(node)
            
        # select all lego parts
        bpy.ops.object.select_all(action='SELECT')
        # make a mirror transform
        bpy.ops.transform.mirror(orient_type='GLOBAL', constraint_axis=(True, False, False))
        time.sleep(0.1)
        # make a rotate about x_axis,with a value of 3.14159(π)
        bpy.ops.transform.rotate(value=3.14159, orient_axis='X', orient_type='GLOBAL', \
                                 orient_matrix=((1, 0, 0), (0, 1, 0), (0, 0, 1)), \
                                 orient_matrix_type='GLOBAL', \
                                 constraint_axis=(True, False, False), \
                                 mirror=True, use_proportional_edit=False, \
                                 proportional_edit_falloff='SMOOTH', \
                                 proportional_size=1, \
                                 use_proportional_connected=False, \
                                 use_proportional_projected=False)
        # unselect all lego parts
        bpy.ops.object.select_all(action = "DESELECT")
        time.sleep(0.1)
        # correct lego parts colors in upper dict
        for i in bpy.data.materials:
            colorName = i.name.split('.')[0]
            mycolors = lego_color_dict.get(colorName)
            if type(mycolors) == list:
                i.use_nodes = False
                for j in range(4):
                    i.diffuse_color[j] = mycolors[j]
        return {'FINISHED'}


class WM_OT_ReadYourColors(Operator):
    bl_label = "read your colors"
    bl_idname = "wm.read_your_colors"
    
    def execute(self, context):
        scene = context.scene
        mytool = scene.my_tool
        color_dict = {}
        with open(mytool.my_path + 'lego_colors.txt', 'r') as f:
            for line in f:
                if line.find(':') > 0:
                    cName, cList = line.split(':')
                    numStr = cList.split('[')[1].split(']')[0]
                    strList = numStr.split(',')
                    numList = []
                    for i in range(len(strList)):
                        numList.append(float(strList[i]))

                    color_dict[cName.split("'")[1].split("'")[0]] = numList
                
        for i in bpy.data.materials:
            colorName = i.name.split('.')[0]
            mycolors = color_dict.get(colorName)
            if type(mycolors) == list:
                i.use_nodes = False
                for j in range(4):
                    i.diffuse_color[j] = mycolors[j]
                    
        return {'FINISHED'}
    
class WM_OT_ParentClearAndKeep(Operator):
    bl_label = "clear parent and keep transform"
    bl_idname = "wm.parent_clear_and_keep"
    
    def execute(self, context):
        scene = context.scene
        mytool = scene.my_tool
        
        # select all lego parts
        bpy.ops.object.select_all(action = "SELECT")
        # clear parent and keep transform
        bpy.ops.object.parent_clear(type='CLEAR_KEEP_TRANSFORM')
        # unselect all lego parts
        bpy.ops.object.select_all(action = "DESELECT")
        
        # delete all parent nodes
        for i in bpy.data.objects:
            if i.name.startswith("Inner-Node") or i.name.startswith("Node"):
                bpy.data.objects.remove(i)

                    
        return {'FINISHED'}

class WM_OT_RemoveLegoLego(Operator):
    bl_label = "remove lego logo"
    bl_idname = "wm.remove_lego_logo"
    
    def execute(self, context):
        scene = context.scene
        mytool = scene.my_tool
        
        for i in bpy.data.meshes:
            if "logo-lego" in i.name:
                bpy.data.meshes.remove(i)
                    
        return {'FINISHED'}
# -----------------------------------------------------------------------
#    Panel in Object Mode
# -----------------------------------------------------------------------
class OBJECT_PT_CustomPanel(Panel):
    bl_label = "Lego Correct"
    bl_idname = "LEGO_Color_change_panel"
    bl_space_type = "VIEW_3D"   
    bl_region_type = "UI"
    bl_category = "LegoCorrect"
    bl_context = "objectmode"   


    @classmethod
    def poll(self,context):
        return context.object is not None

    def draw(self, context):
        layout = self.layout
        scene = context.scene
        mytool = scene.my_tool
        
        layout.operator("wm.lego_dae_correct")
        layout.separator()
        layout.prop(mytool, "my_path")
        layout.operator("wm.read_your_colors")
        layout.separator()
        layout.operator("wm.parent_clear_and_keep")
        layout.operator("wm.remove_lego_logo")
        layout.separator()

# ------------------------------------------------------------------------
#    Registration
# ------------------------------------------------------------------------

classes = (
    MyProperties,
    WM_OT_LegoDaeCorrect,
    WM_OT_ReadYourColors,
    WM_OT_ParentClearAndKeep,
    WM_OT_RemoveLegoLego,
    OBJECT_PT_CustomPanel
)

def register():
    from bpy.utils import register_class
    for cls in classes:
        register_class(cls)

    bpy.types.Scene.my_tool = PointerProperty(type=MyProperties)

def unregister():
    from bpy.utils import unregister_class
    for cls in reversed(classes):
        unregister_class(cls)
    del bpy.types.Scene.my_tool


if __name__ == "__main__":
    register()